import matplotlib.pyplot as plt
import os
import re
import shutil
import string
import tensorflow as tf

from tensorflow.keras import layers
from tensorflow.keras import losses
print(tf.__version__)
url = "https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"

dataset = tf.keras.utils.get_file("aclImdb_v1", url,
                                    untar=True, cache_dir='.',
                                    cache_subdir='')

dataset_dir = os.path.join(os.path.dirname(dataset), 'aclImdb')
batch_size = 32
seed = 42

raw_train_ds = tf.keras.utils.text_dataset_from_directory(
    'aclImdb/train', 
    batch_size=batch_size, 
    validation_split=0.2, 
    subset='training', 
    seed=seed)
    for text_batch, label_batch in raw_train_ds.take(1):
  for i in range(3):
    print("Review", text_batch.numpy()[i])
    print("Label", label_batch.numpy()[i])
    print("Label 0 corresponds to", raw_train_ds.class_names[0])
print("Label 1 corresponds to", raw_train_ds.class_names[1])
raw_val_ds = tf.keras.utils.text_dataset_from_directory(
    'aclImdb/train', 
    batch_size=batch_size, 
    validation_split=0.2, 
    subset='validation', 
    seed=seed)
    raw_test_ds = tf.keras.utils.text_dataset_from_directory(
    'aclImdb/test', 
    batch_size=batch_size)
    def custom_standardization(input_data):
  lowercase = tf.strings.lower(input_data)
  stripped_html = tf.strings.regex_replace(lowercase, '<br />', ' ')
  return tf.strings.regex_replace(stripped_html,
                                  '[%s]' % re.escape(string.punctuation),
                                  '')

                                  max_features = 10000
sequence_length = 250

vectorize_layer = layers.TextVectorization(
    standardize=custom_standardization,
    max_tokens=max_features,
    output_mode='int',
    output_sequence_length=sequence_length)
    def vectorize_text(text, label):
  text = tf.expand_dims(text, -1)
  return vectorize_layer(text), label# retrieve a batch (of 32 reviews and labels) from the dataset
text_batch, label_batch = next(iter(raw_train_ds))
first_review, first_label = text_batch[0], label_batch[0]
print("Review", first_review)
print("Label", raw_train_ds.class_names[first_label])
print("Vectorized review", vectorize_text(first_review, first_label))
model = tf.keras.Sequential([
  layers.Embedding(max_features + 1, embedding_dim),
  layers.Dropout(0.2),
  layers.GlobalAveragePooling1D(),
  layers.Dropout(0.2),
  layers.Dense(1)])

model.summary()
model.compile(loss=losses.BinaryCrossentropy(from_logits=True),
              optimizer='adam',
              metrics=tf.metrics.BinaryAccuracy(threshold=0.0))
              epochs = 10
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs)
    loss, accuracy = model.evaluate(test_ds)

print("Loss: ", loss)
print("Accuracy: ", accuracy)
acc = history_dict['binary_accuracy']
val_acc = history_dict['val_binary_accuracy']
loss = history_dict['loss']
val_loss = history_dict['val_loss']

epochs = range(1, len(acc) + 1)

# "bo" is for "blue dot"
plt.plot(epochs, loss, 'bo', label='Training loss')
# b is for "solid blue line"
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.show()
plt.plot(epochs, acc, 'bo', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend(loc='lower right')

plt.show()
export_model = tf.keras.Sequential([
  vectorize_layer,
  model,
  layers.Activation('sigmoid')
])

export_model.compile(
    loss=losses.BinaryCrossentropy(from_logits=False), optimizer="adam", metrics=['accuracy']
)

# Test it with `raw_test_ds`, which yields raw strings
loss, accuracy = export_model.evaluate(raw_test_ds)
print(accuracy)
examples = [
  "The movie was great!",
  "The movie was okay.",
  "The movie was terrible..."
]

export_model.predict(examples)
import pandas as pd
from scipy.special import expit  # Sigmoid函数

def explain_prediction(text, model, vectorize_layer, top_n=5):
    """解释模型预测，找出影响情感的关键词"""
    # 向量化输入文本
    vectorized_text = vectorize_layer([text])
    tokens = vectorized_text.numpy()[0]
    
    # 获取模型嵌入层和输出层权重
    embedding_layer = model.layers[0]
    dense_layer = model.layers[-1]
    
    # 获取词嵌入
    embeddings = embedding_layer(vectorized_text)
    embeddings = embeddings.numpy()[0]  # shape: (sequence_length, embedding_dim)
    
    # 获取输出层权重（影响情感的方向）
    output_weights = dense_layer.get_weights()[0].flatten()  # shape: (embedding_dim,)
    
    # 计算每个词对情感的贡献度
    contributions = np.dot(embeddings, output_weights)
    
    # 找出贡献最大的词汇
    token_indices = np.argsort(np.abs(contributions))[-top_n:][::-1]
    
    # 转换为单词
    vocab = vectorize_layer.get_vocabulary()
    words = [vocab[tokens[i]] for i in token_indices if tokens[i] != 0]
    scores = [contributions[i] for i in token_indices if tokens[i] != 0]
    
    # 预测结果
    prediction = model.predict(vectorized_text, verbose=0)[0][0]
    sentiment = "积极" if prediction > 0 else "消极"
    
    result = {
        "文本": text,
        "预测情感": sentiment,
        "置信度": expit(prediction),  # 将logits转换为概率
        "关键影响词": list(zip(words, scores))
    }
    return result

# 示例
explanation = explain_prediction("This movie is absolutely amazing, but the plot is a bit predictable.", 
                                model, vectorize_layer)
print("预测解释:", explanation)
# 修改模型结构为LSTM
lstm_model = tf.keras.Sequential([
    layers.Embedding(max_features + 1, embedding_dim),
    layers.Bidirectional(layers.LSTM(64, return_sequences=True)),
    layers.GlobalMaxPool1D(),
    layers.Dropout(0.2),
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.2),
    layers.Dense(1)
])

lstm_model.compile(loss=losses.BinaryCrossentropy(from_logits=True),
                  optimizer='adam',
                  metrics=tf.metrics.BinaryAccuracy(threshold=0.0))

# 训练模型
lstm_history = lstm_model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=5  # LSTM训练较慢，减少epochs
)

# 评估模型
lstm_loss, lstm_accuracy = lstm_model.evaluate(test_ds)
print(f"LSTM模型准确率: {lstm_accuracy:.4f}")
    